#include "kernel_operator.h"
#include "isnan_custom_tiling.h"
#include <cstdint>
#include <cmath>
#include <type_traits>

constexpr int32_t BUFFER_NUM = 2; // tensor num for each queue

constexpr uint32_t BFLOAT16 = 0;            // 对应 DT_BF16，bf16 类型
constexpr uint32_t FLOAT16 = 1;             // 对应 DT_FLOAT16，half（float16）类型
constexpr uint32_t FLOAT32 = 2;             // 对应 DT_FLOAT，float32 类型
constexpr uint32_t INT8 = 3;                // 对应 DT_INT8，int8 类型
constexpr uint32_t INT16 = 4;               // 对应 DT_INT16，int16 类型
constexpr uint32_t INT32 = 5;               // 对应 DT_INT32，int32 类型
constexpr uint32_t UINT8 = 6;               // 对应 DT_UINT8，u8（无符号int8）类型
constexpr uint32_t UINT16 = 7;              // 对应 DT_UINT16，u16（无符号int16）类型
constexpr uint32_t UINT32 = 8;              // 对应 DT_UINT32，u32（无符号int32）类型
constexpr uint32_t INT64 = 9;               // 对应 DT_INT64，int64 类型
constexpr uint32_t UINT64 = 10;             // 对应 DT_UINT64，u64（无符号int64）类型
constexpr uint32_t BOOL = 11;               // 对应 DT_BOOL，bool 类型
constexpr uint32_t DOUBLE = 12;             // 对应 DT_DOUBLE，double（float64）类型
constexpr uint32_t STRING = 13;             // 对应 DT_STRING，string（字符串）类型
constexpr uint32_t DUAL_SUB_INT8 = 14;      // 对应 DT_DUAL_SUB_INT8，dual output int8（双输出int8）类型
constexpr uint32_t DUAL_SUB_UINT8 = 15;     // 对应 DT_DUAL_SUB_UINT8，dual output u8（双输出无符号int8）类型
constexpr uint32_t COMPLEX64 = 16;          // 对应 DT_COMPLEX64，complex64（复数64位）类型
constexpr uint32_t COMPLEX128 = 17;         // 对应 DT_COMPLEX128，complex128（复数128位）类型
constexpr uint32_t QINT8 = 18;              // 对应 DT_QINT8，qint8（量化int8）类型
constexpr uint32_t QINT16 = 19;             // 对应 DT_QINT16，qint16（量化int16）类型
constexpr uint32_t QINT32 = 20;             // 对应 DT_QINT32，qint32（量化int32）类型
constexpr uint32_t QUINT8 = 21;             // 对应 DT_QUINT8，quint8（量化无符号int8）类型
constexpr uint32_t QUINT16 = 22;            // 对应 DT_QUINT16，quint16（量化无符号int16）类型
constexpr uint32_t RESOURCE = 23;           // 对应 DT_RESOURCE，resource（资源）类型
constexpr uint32_t STRING_REF = 24;         // 对应 DT_STRING_REF，string ref（字符串引用）类型
constexpr uint32_t DUAL = 25;               // 对应 DT_DUAL，dual output（双输出）类型
constexpr uint32_t VARIANT = 26;            // 对应 DT_VARIANT，dt_variant（变体）类型
constexpr uint32_t UNDEFINED = 27;          // 对应 DT_UNDEFINED，未设置的DataType字段
constexpr uint32_t INT4 = 28;               // 对应 DT_INT4，int4 类型
constexpr uint32_t UINT1 = 29;              // 对应 DT_UINT1，u1（无符号int1）类型
constexpr uint32_t INT2 = 30;               // 对应 DT_INT2，int2 类型
constexpr uint32_t UINT2 = 31;              // 对应 DT_UINT2，u2（无符号int2）类型
constexpr uint32_t COMPLEX32 = 32;          // 对应 DT_COMPLEX32，complex32（复数32位）类型
constexpr uint32_t HIFLOAT8 = 33;           // 对应 DT_HIFLOAT8，hif8（高精度float8）类型
constexpr uint32_t FLOAT8_E5M2 = 34;        // 对应 DT_FLOAT8_E5M2，fp8_e5m2（float8格式，5位指数2位尾数）类型
constexpr uint32_t FLOAT8_E4M3FN = 35;      // 对应 DT_FLOAT8_E4M3FN，fp8_e4m3（float8格式，4位指数3位尾数）类型
constexpr uint32_t FLOAT8_E8M0 = 36;        // 对应 DT_FLOAT8_E8M0，fp8_e8m0（float8格式，8位指数0位尾数）类型
constexpr uint32_t FLOAT6_E3M2 = 37;        // 对应 DT_FLOAT6_E3M2，fp6_e3m2（float6格式，3位指数2位尾数）类型
constexpr uint32_t FLOAT6_E2M3 = 38;        // 对应 DT_FLOAT6_E2M3，fp6_e2m3（float6格式，2位指数3位尾数）类型
constexpr uint32_t FLOAT4_E2M1 = 39;        // 对应 DT_FLOAT4_E2M1，fp4_e2m1（float4格式，2位指数1位尾数）类型
constexpr uint32_t FLOAT4_E1M2 = 40;        // 对应 DT_FLOAT4_E1M2，fp4_e1m2（float4格式，1位指数2位尾数）类型
constexpr uint32_t MAX = 41;                // 对应 DT_MAX，标记AscendC数据类型的边界

constexpr uint32_t LAST_TWO_TILE = 2;

// 定义复数类型
struct complex64 {
    float real;
    float imag;
};

// 定义复数类型
struct complex128 {
    double real;
    double imag;
};

__aicore__ inline bool IsFloat32NaN(float val) {
    return (*reinterpret_cast<const uint32_t*>(&val) == 0x7FC00000);
}

__aicore__ inline bool IsDoubleNaN(double val) {
    return (*reinterpret_cast<const uint64_t*>(&val) == 0x7FF8000000000000);
}

__aicore__ inline bool IsHalfNaN(half val) {
    return (*reinterpret_cast<const uint16_t*>(&val) == 0x7E00);
}

__aicore__ inline bool IsBfloat16NaN(bfloat16_t val) {
    return (*reinterpret_cast<const uint16_t*>(&val) == 0x7F80);
}

__aicore__ inline bool IsComplex64NaN(float real, float imag) {
    return IsFloat32NaN(real) || IsFloat32NaN(imag);
}

__aicore__ inline bool IsComplex128NaN(double real, double imag) {
    return IsDoubleNaN(real) || IsDoubleNaN(imag);
}

template <typename dataType>
class KernelIsnan {
public:
    __aicore__ inline KernelIsnan() {}

    __aicore__ inline void Init(GM_ADDR x, GM_ADDR out, IsnanCustomTilingData tiling)
    {
        this->dataTypeId = tiling.dataType;
        if (tiling.isEvenCore)
        {
            this->blockLength = tiling.blockLength;
            this->tileNum = tiling.tileNum;
            this->tileLength = tiling.tileLength / BUFFER_NUM;
            this->lastTileLength = tiling.lastTileLength;

            xGm.SetGlobalBuffer((__gm__ dataType *)x + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
            outGm.SetGlobalBuffer((__gm__ bool *)out + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
        }
        else
        {
            if (AscendC::GetBlockIdx() < tiling.formerNum)
            {
                this->tileNum = tiling.formerTileNum;
                this->tileLength = tiling.formerTileLength / BUFFER_NUM;
                this->lastTileLength = tiling.formerLastTileLength;

                xGm.SetGlobalBuffer((__gm__ dataType *)x + tiling.formerLength * AscendC::GetBlockIdx(),
                                    tiling.formerLength);
                outGm.SetGlobalBuffer((__gm__ bool *)out + tiling.formerLength * AscendC::GetBlockIdx(),
                                    tiling.formerLength);
            }
            else
            {
                this->tileNum = tiling.tailTileNum;
                this->tileLength = tiling.tailTileLength / BUFFER_NUM;
                this->lastTileLength = tiling.tailLastTileLength;

                xGm.SetGlobalBuffer((__gm__ dataType *)x + tiling.formerLength * tiling.formerNum +
                                        tiling.tailLength * (AscendC::GetBlockIdx() - tiling.formerNum),
                                    tiling.tailLength);
                outGm.SetGlobalBuffer((__gm__ bool *)out + tiling.formerLength * tiling.formerNum +
                                        tiling.tailLength * (AscendC::GetBlockIdx() - tiling.formerNum),
                                    tiling.tailLength);
            }
        }
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(dataType));
        pipe.InitBuffer(outQueueBool, BUFFER_NUM, this->tileLength * sizeof(bool));
    }

    __aicore__ inline void Process() {
        int32_t loopCount = this->tileNum * BUFFER_NUM;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t progress) {
        AscendC::LocalTensor<dataType> xLocal = inQueueX.AllocTensor<dataType>();
        if ((progress == (this->tileNum * BUFFER_NUM - 2)) || (progress == (this->tileNum * BUFFER_NUM - 1)))
        {
            AscendC::DataCopy(xLocal, xGm[(progress - LAST_TWO_TILE) * this->tileLength + this->lastTileLength],
                              this->tileLength);
        }
        else
        {
            AscendC::DataCopy(xLocal, xGm[progress * this->tileLength], this->tileLength);
        }
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void Compute(int32_t progress) {
        AscendC::LocalTensor<dataType> xLocal = inQueueX.DeQue<dataType>();
        AscendC::LocalTensor<bool> outLocal = outQueueBool.AllocTensor<bool>();
        bool isNaN = false;

        for (int32_t i = 0; i < this->tileLength; i++) {
            // 编译期判断模板类型，只保留当前类型的代码
            if constexpr (std::is_same_v<dataType, float>) {
                // float32 分支：正常用 GetValue()
                float element = xLocal.GetValue(i);
                isNaN = IsFloat32NaN(element);
            } else if constexpr (std::is_same_v<dataType, double>) {
                // double 分支
                double element = xLocal.GetValue(i);
                isNaN = IsDoubleNaN(element);
            } else if constexpr (std::is_same_v<dataType, half>) {
                // float16 分支
                half element = xLocal.GetValue(i);
                isNaN = IsHalfNaN(element);
            } else if constexpr (std::is_same_v<dataType, bfloat16_t>) {
                // bfloat16 分支
                bfloat16_t element = xLocal.GetValue(i);
                isNaN = IsBfloat16NaN(element);
            } else if constexpr (std::is_same_v<dataType, complex64>) {
                // AscendC::LocalTensor 没法直接返回 complex 类型，这里用指针访问
                // auto* ptr = reinterpret_cast<complex64*>(xLocal.GetPhyAddr());
                // complex64 element = ptr[i];
                AscendC::LocalTensor<float> xLocal_rc = xLocal.template ReinterpretCast<float>();
                float real = xLocal_rc.GetValue(2 * i);
                float imag = xLocal_rc.GetValue(2 * i + 1);
                isNaN = IsComplex64NaN(real, imag);
            } else if constexpr (std::is_same_v<dataType, complex128>) {
                // auto* ptr = reinterpret_cast<complex128*>(xLocal.GetPhyAddr());
                // complex128 element = ptr[i];
                AscendC::LocalTensor<double> xLocal_rc = xLocal.template ReinterpretCast<double>();
                double real = xLocal_rc.GetValue(2 * i);
                double imag = xLocal_rc.GetValue(2 * i + 1);
                isNaN = IsComplex128NaN(real, imag);
            } else {
                // 整数/布尔类型：无 NaN，直接返回 false
                isNaN = false;
            }

            outLocal.SetValue(i, isNaN);
        }

        outQueueBool.EnQue(outLocal);
        inQueueX.FreeTensor(xLocal);
    }

    __aicore__ inline void CopyOut(int32_t progress) {
        AscendC::LocalTensor<bool> outLocal = outQueueBool.DeQue<bool>();
        AscendC::DataCopyExtParams copyParams{1, (uint32_t)(this->tileLength * sizeof(bool)), 0, 0, 0};

        if((progress == (this->tileNum * BUFFER_NUM - 2)) || (progress == (this->tileNum * BUFFER_NUM - 1))){
            AscendC::DataCopyPad(outGm[(progress - LAST_TWO_TILE) * this->tileLength + this->lastTileLength], outLocal, copyParams);
        }else{
            AscendC::DataCopyPad(outGm[progress * this->tileLength], outLocal, copyParams);
        }
        outQueueBool.FreeTensor(outLocal);
    }
    
private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueBool;
    AscendC::GlobalTensor<dataType> xGm;
    AscendC::GlobalTensor<bool> outGm;
    AscendC::DataCopyExtParams copyParams;

    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t lastTileLength;

    uint32_t dataTypeId;
};


extern "C" __global__ __aicore__ void isnan_custom(GM_ADDR x, GM_ADDR out, IsnanCustomTilingData tiling)
{
    if (tiling.dataType == BFLOAT16) {
        KernelIsnan<bfloat16_t> op;
        op.Init(x, out, tiling);
        op.Process();
    } else if (tiling.dataType == FLOAT16) {
        KernelIsnan<half> op;
        op.Init(x, out, tiling);
        op.Process();
    } else if (tiling.dataType == FLOAT32) {
        KernelIsnan<float> op;
        op.Init(x, out, tiling);
        op.Process();
    } else if (tiling.dataType == DOUBLE) {
        KernelIsnan<double> op;
        op.Init(x, out, tiling);
        op.Process();
    } else if (tiling.dataType == INT8) {
        KernelIsnan<int8_t> op;
        op.Init(x, out, tiling);
        op.Process();
    } else if (tiling.dataType == INT16) {
        KernelIsnan<int16_t> op;
        op.Init(x, out, tiling);
        op.Process();
    } else if (tiling.dataType == INT32) {
        KernelIsnan<int32_t> op;
        op.Init(x, out, tiling);
        op.Process();
    }else if (tiling.dataType == COMPLEX64) {
        KernelIsnan<complex64> op;
        op.Init(x, out, tiling);
        op.Process();
    }else if (tiling.dataType == COMPLEX128) {
        KernelIsnan<complex128> op;
        op.Init(x, out, tiling);
        op.Process();
    } 
    else {
        return;
    }
}
